import sys

from attribute_chooser import AttributeChooser
from decision_tree import DecisionTree

class DecisionTreeClassifier():
  def __init__(self, input_filename,
      attribute_chooser_method = 'information_gain'):

    # Attributes
    attributes = ('STABILITY', 'ERROR', 'SIGN', 'WIND',
        'MAGNITUDE', 'VISIBILITY')

    # Concepts
    concepts = ('noauto', 'auto')

    # Attribute value dictionary
    attribute_value_dict = {}   
    attribute_value_dict['STABILITY'] = ('stab', 'xstab')
    attribute_value_dict['ERROR'] = ('XL', 'LX', 'MM', 'SS')
    attribute_value_dict['SIGN'] = ('pp', 'nn')
    attribute_value_dict['WIND'] = ('head', 'tail')
    attribute_value_dict['MAGNITUDE'] = ('low', 'medium', 'strong',
    'outofrange')
    attribute_value_dict['VISIBILITY'] = ('yes', 'no')

    # Read the data file and create the feature vectors.
    in_file = open(input_filename)
    feature_vectors = []
    for line in in_file:
      line = line.strip()
      parts = line.split(' ')
      label = concepts[int(parts[0]) - 1]
      features = {}
      for idx in range(1, len(parts)):
        attribute = attributes[idx - 1]
        attribute_value_idx = int(parts[idx]) - 1
        attribute_value = attribute_value_dict[attribute][attribute_value_idx]
        features[attribute] = attribute_value
      feature_vectors.append((features, label))
    in_file.close()

    # Create the decision tree based on the attribute_chooser
    attribute_chooser = AttributeChooser(attribute_chooser_method)
    decision_tree = DecisionTree(feature_vectors, attribute_value_dict,
        'noauto', attribute_chooser)
    decision_tree.print_tree()

attribute_chooser_method = 'information_gain'
if len(sys.argv) == 2:
  attribute_chooser_method = sys.argv[1]
classifier = DecisionTreeClassifier('./shuttle_ext_unique.dat',
    attribute_chooser_method)
